const fs = require('fs').promises;
const path = require('path');
const BaseDatasetWrapper = require('../../BaseDatasetWrapper');
const COCOLabelLoader = require('./COCOLabelLoader');
const { DatasetListDB } = require('../../../dbs/dbDatasetList');
const { getAllFilesInDirectory } = require('../../../utils/fileUtils');
const { getImageSize } = require('../../../utils/imageUtils');

class COCODatasetWrapper extends BaseDatasetWrapper {
    // 读入目录信息，解析对应的标注结果
    async parse(dataset_info) {
        // console.log('COCODatasetWrapper::parse dataset_info = ', dataset_info);
        const datasetId = dataset_info.dataset_id;
        const datasetDir = dataset_info.dataset_dir;
        this.parseCallback(datasetId, 'parse_status', '读取图片列表...');
        let imageFiles = await getAllFilesInDirectory(datasetDir, ['.jpg', '.jpeg', '.png', '.bmp'], ['workspace']);
        this.parseCallback(datasetId, 'parse_start', imageFiles.length);

        // 判断是否有selections(打开某一个coco json标注文件)
        let jsonFiles = await getAllFilesInDirectory(datasetDir, ['.json'], ['workspace']);
        const selections = JSON.parse(dataset_info.label_selections);
        if (selections.length > 0) {
            jsonFiles = selections.map((file) => {
                return path.join(datasetDir, file);
            });
        }

        // 创建图像文件路径的映射关系
        let totalSamples = imageFiles.length;
        let matchingPairs = [];
        let categories = null;
        let imageFilePathMap = {};
        imageFiles.forEach((imagePath) => {
            const imageName = path.basename(imagePath);
            imageFilePathMap[imageName] = {
                path: imagePath,
                labeled: false
            };
        });

        // 创建work目录结构
        const workDir = path.join(datasetDir, 'workspace', datasetId);
        await fs.mkdir(workDir, { recursive: true });

        // 对每个JSON文件进行处理
        for (let jsonFilePath of jsonFiles) {
            // console.log('parse JsonFile = ', jsonFilePath);
            this.parseCallback(datasetId, 'parse_status', `解析${path.basename(jsonFilePath)}文件...`);

            let annotationMap = {};
            try {
                // 读取jsonFile
                const jsonData = await fs.readFile(jsonFilePath, 'utf8');
                const cocoJson = JSON.parse(jsonData);

                // 检查JSON文件是否符合COCO格式
                if (!cocoJson.images || !cocoJson.annotations) {
                    console.warn(`The file ${jsonFilePath} does not appear to be in COCO format.`);
                    continue;
                }

                // 处理并记录categories信息
                if (cocoJson.categories) {
                    categories = {};
                    cocoJson.categories.forEach((category) => {
                        categories[category.id] = category;
                    });
                }

                // 创建用于存储每张图像对应的标注信息的映射关系
                cocoJson.annotations.forEach((annotation) => {
                    if (categories) {
                        let category_id = annotation['category_id'];
                        annotation['category'] = categories[category_id].name;
                    }

                    const imageId = annotation.image_id;
                    if (!annotationMap[imageId]) {
                        annotationMap[imageId] = [];
                    }
                    annotationMap[imageId].push(annotation);
                });
                // console.log('annotationMap = ', annotationMap);

                // 处理每张图像信息
                cocoJson.images.forEach((image) => {
                    const imageName = path.basename(image.file_name);
                    const imageId = image.id;
                    const imagePathObj = imageFilePathMap[imageName];
                    // console.log(`imageName = ${imageName}, imagePathObj = `, imagePathObj);
                    // console.log(`imageId = ${imageId}, annotationMap[imageId] = `, annotationMap[imageId]);
                    if (imagePathObj) {
                        imagePathObj.labeled = true;
                        imagePathObj.annotations = annotationMap[imageId] || [];
                    } else {
                        // console.warn('Unfound image name =', imageName);
                    }
                });
            } catch (error) {
                console.error(`Error reading or parsing JSON file ${jsonFilePath}: `, error);
            }
        }

        // 统一处理所有图像（包括已标注和未标注）
        let i = 0;
        let labeledSamples = 0;
        let imageCount = Object.keys(imageFilePathMap).length;
        // console.log('imageFilePathMap len = ', imageCount);
        for (const imageName in imageFilePathMap) {
            const imagePathObj = imageFilePathMap[imageName];
            const imagePath = imagePathObj.path;

            // 根据图片的相对路径，创建需要复制到的目标目录结构
            const [labelFileDir, labelFilePath] = this.getLabelPath(datasetDir, workDir, imageName, imagePath);

            // 确保目标目录存在
            await fs.mkdir(labelFileDir, { recursive: true });

            // 获取标注信息（已标注图像从 annotations 中获取，未标注图像为空数组）
            let imageAnnotations = [];
            if (imagePathObj.labeled) {
                imageAnnotations = imagePathObj.annotations;
                labeledSamples += 1;
            }

            // 添加到 matchingPairs
            matchingPairs.push({
                image_path: path.relative(datasetDir, imagePath),
                label_path: path.relative(datasetDir, labelFilePath)
            });

            // 写入标注信息
            await fs.writeFile(labelFilePath, JSON.stringify(imageAnnotations));

            // 更新数据加载进展
            this.parseCallback(datasetId, 'parse_percent', (((i + 1) * 100) / imageCount).toFixed(2));
            i += 1;
        }

        // 保存matching pairs结果
        await this.saveMatchingPairs(datasetId, matchingPairs);

        // 汇总图片目录名称，方便后续切分训练集和验证集使用
        let imageDirs = new Set();
        imageFiles.forEach((file) => imageDirs.add(path.dirname(file)));
        await this.createSubsets(dataset_info, Array.from(imageDirs));

        // 保存categories信息
        // console.log('categories = ', categories);
        // console.log(dataset_info);
        if (categories) {
            this.saveCategories(dataset_info.dataset_id, categories);
        }

        return {
            code: 0,
            msg: {
                totalSamples: totalSamples,
                labeledSamples: labeledSamples
            }
        };
    }

    getLabelPath(datasetDir, workDir, imageName, imageFilePath) {
        const imageRelativePath = path.relative(datasetDir, imageFilePath);
        const labelRelativeDir = path.dirname(imageRelativePath);
        const newLabelFileDir = path.join(workDir, labelRelativeDir);
        const jsonFileName = path.basename(imageName, path.extname(imageName)) + '.json';
        const imageAnnotationsPath = path.join(newLabelFileDir, jsonFileName);
        return [newLabelFileDir, imageAnnotationsPath];
    }

    // 导出数据集标注结果列表
    async export(toFormat, dataset_info, subsets_info) {
        // console.log('COCODatasetWrapper::export dataset_info = ', dataset_info);
        // console.log('COCODatasetWrapper::export subsets_info = ', subsets_info);
        // 判断dataset是否是COCO数据集，从而从不同位置读取label
        let label_format = dataset_info.label_format.toLowerCase();

        // 将train、val、test的子集目录，以及子集的图片文件列表，写入到本地文件中
        let dataset_id = dataset_info.id;
        let dataset_dir = dataset_info.dataset_dir;
        let workDir = path.join(dataset_dir, 'workspace', dataset_info.id);
        for (let i = 0; i < subsets_info.length; i++) {
            // 输出的coco json
            let cocoJson = {
                images: [],
                annotations: [],
                categories: []
            };

            // 读取子集的图片文件列表
            let subset = subsets_info[i];
            let imageDir = path.join(dataset_dir, subset.image_dir);
            let labelDir = path.join(dataset_dir, subset.label_dir);
            let imageFiles = await getAllFilesInDirectory(imageDir, ['.jpg', '.jpeg', '.png', '.bmp'], ['workspace']);
            // console.log('subset = ', subset);
            // console.log('imageDir = ', imageDir);

            // 遍历每一个图片文件，合并标注结果到cocoJson中
            let anno_id = 1; // COCO数据集从1开始
            for (let j = 0; j < imageFiles.length; j++) {
                let imagePath = imageFiles[j];
                let baseName = path.parse(imagePath).name;
                let imageName = path.basename(imagePath);

                let tmpPath = path.join('workspace', `${dataset_id}`, path.relative(dataset_dir, imagePath));
                if (label_format !== 'coco') {
                    tmpPath = path.join('workspace', `${dataset_id}`, 'converted_coco', path.relative(dataset_dir, imagePath));
                }

                const parsedPath = path.parse(tmpPath);
                const newLabelPath = path.format({
                    ...parsedPath,
                    base: undefined,
                    ext: '.json'
                });
                const labelPath = path.join(dataset_dir, newLabelPath);
                // console.log('COCODatasetWrapper::export labelPath = ', labelPath);

                // 保存图像信息
                let size = getImageSize(imagePath);
                // console.log(size);
                cocoJson.images.push({
                    id: j,
                    width: size.width,
                    height: size.height,
                    file_name: imageName
                });

                // 保存标注信息
                const jsonData = await fs.readFile(labelPath, 'utf8');
                const jsonAnnos = JSON.parse(jsonData);
                for (let k = 0; k < jsonAnnos.length; k++) {
                    let jsonAnno = jsonAnnos[k];
                    jsonAnno.image_id = j;
                    jsonAnno.id = anno_id;
                    anno_id++;
                }
                // console.log('jsonAnnos = ', jsonAnnos);
                cocoJson.annotations = cocoJson.annotations.concat(jsonAnnos);
            }

            // 保存类别信息
            const categories = await this.app.db.categories.getCategories(dataset_id);
            for (let c = 0; c < categories.length; c++) {
                let category = categories[c];
                let jsonCat = {
                    id: category.category_id,
                    name: category.category_name
                };
                if (category.supercategory) {
                    jsonCat['supercategory'] = category.supercategory;
                }
                cocoJson.categories.push(jsonCat);
            }

            // 保存json结果
            const jsonString = JSON.stringify(cocoJson);
            const output_file = path.join(workDir, `${subset.subset_type}.json`);
            await fs.writeFile(output_file, jsonString, 'utf-8');
            subset.label_list_path = path.relative(dataset_dir, output_file);
        }
    }

    async convert(toFormat, dataset_info) {
        let format = toFormat.toLowerCase();
        if (format === 'coco') {
            return true;
        }
        const datasetId = dataset_info.id;
        const datasetDir = dataset_info.dataset_dir;

        try {
            // 获取数据集的标注列表
            let dblist = this.app.list_dbs[datasetId];
            if (!dblist) {
                this.app.list_dbs[datasetId] = new DatasetListDB(this.app, datasetId);
                dblist = this.app.list_dbs[datasetId];
            }
            let pairs = await dblist.pairs.getPairs();

            // 动态加载labelloader类
            const srcLoader = new COCOLabelLoader(this.app);
            const loaderFile = toFormat.toUpperCase() + 'LabelLoader.js';
            const loaderClass = require(path.join(__dirname, loaderFile));
            const dstloader = new loaderClass(this.app);

            // 对于每一个pair进行转换
            let categories = await this.app.db.categories.getCategories(datasetId);
            for (let i = 0; i < pairs.length; i++) {
                let pair = pairs[i];
                pair.image_path = path.join(datasetDir, pair.image_path);
                pair.label_path = path.join(datasetDir, pair.label_path);
                let label_data = await srcLoader.load({
                    dataset_id: datasetId,
                    image_path: pair.image_path,
                    label_path: pair.label_path
                });
                let label_wrapper = this.wrapLabelData(dataset_info, pair, label_data);
                let label_info = await dstloader.toString(label_wrapper, categories);

                // 把结果保存到本地文件中
                let outputPath = BaseDatasetWrapper.getConvertedPath(toFormat, dataset_info, pair);
                await fs.writeFile(outputPath, label_info);
            }
            return { code: 0, msg: `COCO to ${toFormat.toUpperCase()} successfully` };
        } catch (e) {
            console.error('COCODatasetWrapper::convert error = ', e);
            return { code: -1, msg: e.message };
        }
    }
}

module.exports = COCODatasetWrapper;
